package com.mcy.springbatch.job;

import com.mcy.springbatch.entity.SysUser;
import com.mcy.springbatch.partitioner.RangePartitioner;
import com.mcy.springbatch.reader.ReadUserDataJpaPagingItemReader;
import com.mcy.springbatch.writer.SysUserItemWriter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.core.launch.support.RunIdIncrementer;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.batch.item.database.JpaPagingItemReader;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;
import java.util.HashMap;
import java.util.Map;

/**
 * 分区配置，分区的原理是一个master主处理器对应多个从slave处理器
 * @author : Create By mcy
 * @date : 2019/8/26 17:21:47
 */
@Slf4j
@Configuration
public class RPartitionerJobConfig extends BaseJobConfig {

    @Autowired
    private ReadUserDataJpaPagingItemReader readUserDataJpaPagingItemReader;

    @Autowired
    private SysUserItemWriter sysUserItemWriter;

    @Autowired
    RangePartitioner rangePartitioner;

    @Bean
    public Job PartitionJob() {
        return jobBuilderFactory.get("partitionJob")
                .incrementer(new RunIdIncrementer())
                .start(masterStep())
                .build();
    }

    @Bean
    public Step masterStep() {
        return stepBuilderFactory.get("masterStep")
                .partitioner("slave", rangePartitioner)
                .partitionHandler(masterSlaveHandler())
                .build();
    }


    @Bean
    public PartitionHandler masterSlaveHandler() {
        TaskExecutorPartitionHandler handler = new TaskExecutorPartitionHandler();
        handler.setGridSize(12);
        handler.setTaskExecutor(new SimpleAsyncTaskExecutor());
        handler.setStep(slave());
        try {
            handler.afterPropertiesSet();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return handler;
    }

    @Bean(name = "slave")
    public Step slave() {
        return stepBuilderFactory.get("slave")
                .<SysUser, SysUser>chunk(50)
                .reader(slaveReader(0, 0))
                .writer(sysUserItemWriter)
                .build();
    }

    @Bean
    @StepScope
    public JpaPagingItemReader<SysUser> slaveReader(
            @Value("#{stepExecutionContext[fromId]}")  Integer fromId,
            @Value("#{stepExecutionContext[toId]}")  Integer toId) {
        log.info("slaveReader start " + fromId + " " + toId);
        JpaPagingItemReader<SysUser> reader = new JpaPagingItemReader<>();
        reader.setEntityManagerFactory(entityManager);
        reader.setPageSize(20);
        reader.setQueryString("select new SysUser(id,loginName,userName,idCard,phone,sex,address,email,status) from SysUser where id >= :fromId and id <= :toId");
        Map<String, Object> params = new HashMap<>();
        try {
            params.put("fromId", fromId);
            params.put("toId", toId);
        } catch (Exception e) {

        }
        reader.setParameterValues(params);
        reader.setSaveState(false);
        return reader;
    }

}
